package kelvin.booleanretrieval;

import java.io.IOException;
import java.util.Iterator;
import java.util.StringTokenizer;
import java.util.ArrayList;
import java.util.Collections;
import java.io.DataInput;
import java.io.DataOutput;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.FloatWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.Mapper;
import org.apache.hadoop.mapreduce.Reducer;
import org.apache.hadoop.mapreduce.Partitioner;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.GenericOptionsParser;

public class BooleanRetrieval {

  public static class BooleanRetrievalMapper extends Mapper<LongWritable, Text, Text, IntWritable> {

    private String getClean(String s) {
      String t = "";
      for (int i = 0; i < s.length(); i++)
        if (Character.isLetterOrDigit(s.charAt(i)))
          t += s.charAt(i);
        else
          t += " ";
      return t;
    }

    private String queryWords;

    public void setup(Context context) {
      queryWords = context.getConfiguration().get("query_words");
    }

    static private IntWritable one = new IntWritable(1);

    /* 
     * input: (line-offset, word : #postings : (line1, count1) (line2, count2))
     * output: (line, 1)... for matched word
     */
    public void map(LongWritable lineOffset, Text line, Context context) throws IOException, InterruptedException {
      StringTokenizer itr = new StringTokenizer(getClean(line.toString()));
      String word = itr.nextToken();
      if (queryWords.indexOf(" " + word + " ") == -1)
        return;
      int num = Integer.parseInt(itr.nextToken());
      for (int j = 0; j < num; j++) {
        String offset = itr.nextToken();
        String count = itr.nextToken();
        context.write(new Text(offset), one);
      }
    }
  }

  public static class BooleanRetrievalReducer extends Reducer<Text, IntWritable, Text, Text> {

    private int numQueryWords;

    public void setup(Context context) {
      numQueryWords = context.getConfiguration().getInt("num_query_words", -1);
    }

    static private Text emptyText = new Text("");

		public void reduce(Text word, Iterable<IntWritable> values, Context context) throws IOException, InterruptedException {
      int count = 0;
      for (IntWritable i : values)
        count++;
      if (count == numQueryWords)
        context.write(word, emptyText);
		}
	}
  
  public static void main(String[] args) throws Exception {
    Configuration conf = new Configuration();
    args = new GenericOptionsParser(conf, args).getRemainingArgs();
    if (args.length <= 2) {
      System.err.println("Usage: BooleanRetrieval <input-inverted-index> <output> <word1> <word2> ...");
      System.exit(2);
    }

    String inputDirName = args[0];
    String outputDirName = args[1];
    String queryWords = "";
    for (int i = 2; i < args.length; i++)
      queryWords += " " + args[i] + " ";
    conf.set("query_words", queryWords);
    conf.setInt("num_query_words", args.length - 2);

    System.out.printf("Query words: %s\n", queryWords);

    Job job = new Job(conf, "Boolean Retrieval");
    job.setJarByClass(BooleanRetrieval.class);
    job.setMapperClass(BooleanRetrievalMapper.class);
    job.setReducerClass(BooleanRetrievalReducer.class);
    job.setMapOutputKeyClass(Text.class);
    job.setMapOutputValueClass(IntWritable.class);
    job.setOutputKeyClass(Text.class);
    job.setOutputValueClass(Text.class);
    job.setInputFormatClass(TextInputFormat.class);
    job.setOutputFormatClass(TextOutputFormat.class);
    FileInputFormat.addInputPath(job, new Path(inputDirName));
    FileOutputFormat.setOutputPath(job, new Path(outputDirName));
    System.exit(job.waitForCompletion(true) ? 0 : 1);
  }
}

